package com.google.code.securityflow.cfg

import com.google.code.securityflow.graphviz.GraphvizCluster
import com.google.code.securityflow.ast._
import com.google.code.securityflow.util.MultiMap
import java.lang.RuntimeException

object FlowGraph {

  class Counter {
    var counter = 0;

    def next(): Int = {
      counter = counter + 1;
      counter;
    }
  }

  val functionCounter = new Counter();
  val blockCounter = new Counter();
  val nodeCounter = new Counter();

  val tempVar = new Counter();

  def build(program: Program): FlowGraph = {
    val mainFunction = parseFunction(program.body);

    /**
     * Returns the set of functions declared in the given function.
     * If the given function is main then all functions in the program are returned.
     */
    def findNestedFunctions(function: Function): Set[Function] = {
      function.basicBlocks.flatMap(block => {
        block.nodes.foldLeft(Set.empty[Function])((acc, node) => {
          node match {
            case FunctionLiteralNode(_, f, _, _) => acc + f;
            case _ => acc;
          }
        })
      })
    }

    // All functions declared in the program (including "main")
    val functions = findNestedFunctions(mainFunction) + mainFunction;

    // Add NoOp nodes to empty blocks.
    for (function <- functions; block <- function.basicBlocks; if block.isEmpty) {
      block match {
        case normalBlock: NormalBlock => {
          normalBlock.addNode(NoOpNode(nodeCounter.next(), SourceLocation.synthetic))
        }
        case _ => {
          // do nothing
        }
      }
    }

    // Build auxiliary data structures

    // Functions
    val functionOf = functions.flatMap(f => f.basicBlocks.map(b => (b, f))).toMap;

    // Successors
    var successorsOf = MultiMap.empty[BasicBlock, BasicBlock];
    for (function <- functions; block <- function.basicBlocks) {
      successorsOf = successorsOf ++ (block -> block.successors);
    }

    // Predecessors
    val predecessorsOf = successorsOf.inverse;

    // Descendants
    var descendantsOf = MultiMap.empty[BasicBlock, BasicBlock];
    for (function <- functions) {
      var fixpoint = false;
      for (block <- function.basicBlocks) {
        descendantsOf = descendantsOf ++ (block -> block.successors);
      }
      while (!fixpoint) {
        val prev = descendantsOf;
        for (block <- function.basicBlocks) {
          val curr = descendantsOf(block);
          val next = curr.flatMap(b => b.successors);
          descendantsOf = descendantsOf ++ (block -> next);
        }
        fixpoint = prev == descendantsOf;
      }
    }

    // Ancestors
    val ancestorsOf = descendantsOf.inverse;

    new FlowGraph(mainFunction, functions, functionOf, predecessorsOf, successorsOf, ancestorsOf, descendantsOf);
  }

  def parseFunction(block: Block): Function = {
    parseStatementList(block.statements);
  }

  def parseStatementList(statements: List[Statement]): Function = {
    val firstBlock = new NormalBlock(blockCounter.next());
    var currentBlock = firstBlock;
    for (stm <- statements) {
      currentBlock = parseStatement(stm, currentBlock);
    }
    val entryBlock = EntryBlock(blockCounter.next(), firstBlock);
    val exitBlock = ExitBlock(blockCounter.next());
    currentBlock.addSuccessor(exitBlock);
    new Function(functionCounter.next(), entryBlock, exitBlock);
  }

  def parseStatements(statements: List[Statement], currentBlock: NormalBlock): NormalBlock = {
    var current = currentBlock;
    for (stm <- statements) {
      current = parseStatement(stm, currentBlock);
    }
    current;
  }

  /**
   * Build BasicBlocks for a Statement.
   *
   * @param statement the statement
   * @param currentBlock the basic block to add nodes to
   * @return the next basic block to add nodes to
   */
  def parseStatement(statement: Statement, currentBlock: NormalBlock): NormalBlock = {
    statement match {
      case OutStatement(id, expression, location) => {
        val (nextBlock, valueVar) = parseExpression(expression, currentBlock);
        val node = OutNode(nodeCounter.next(), valueVar, location);
        nextBlock.addNode(node);
        nextBlock
      }
      // Assignment: Case 1: Write to variable
      case AssignStatement(id, SimplePropertyAccessExpression(_, name, _), expression, location) => {
        val (nextBlock, valueVar) = parseExpression(expression, currentBlock);
        val node = WriteVariableNode(nodeCounter.next(), name, valueVar, location);
        nextBlock.addNode(node);
        nextBlock
      }
      // Assignment: Case 2: Write to property
      case AssignStatement(id, QualifiedPropertyAccessExpression(_, propertyAccess, name, _), expression, location) => {
        val baseVar = parseReadPropertyExpression(propertyAccess, currentBlock);
        val (block, valueVar) = parseExpression(expression, currentBlock);
        val node = WritePropertyNode(nodeCounter.next(), name, baseVar, valueVar, location);
        block.addNode(node);
        block
      }
      case CallStatement(id, expression, location) => {
        val (block, resultVar) = parseExpression(expression, currentBlock);

        // Add after-call block
        val afterCallBlock = new NormalBlock(blockCounter.next());

        val node = CallNode(nodeCounter.next(), resultVar, location);
        val callBlock = new CallBlock(blockCounter.next(), node, afterCallBlock);
        block.addSuccessor(callBlock);

        afterCallBlock;
      }
      case IfStatement(id, expression, statements, location) => {
        // Build conditional expression
        val (block, resultVar) = parseExpression(expression, currentBlock);

        val trueBlock = new NormalBlock(blockCounter.next());
        val joinBlock = new NormalBlock(blockCounter.next());

        // Add IfNode
        block.addNode(IfNode(nodeCounter.next(), resultVar, joinBlock, location));

        block.addSuccessor(trueBlock);
        block.addSuccessor(joinBlock);

        val nextBlock = parseStatements(statements, trueBlock);
        nextBlock.addSuccessor(joinBlock);

        joinBlock;
      }
      case WhileStatement(id, expression, statements, location) => {
        val conditionBlock = new NormalBlock(blockCounter.next());
        val loopEntry = new NormalBlock(blockCounter.next());
        val joinBlock = new NormalBlock(blockCounter.next());

        currentBlock.addSuccessor(conditionBlock);
        conditionBlock.addSuccessor(loopEntry);
        conditionBlock.addSuccessor(joinBlock);

        val (condBB, resultVar) = parseExpression(expression, conditionBlock);
        condBB.addNode(IfNode(nodeCounter.next(), resultVar, joinBlock, location));

        val loopExit = parseStatements(statements, loopEntry);

        loopExit.addSuccessor(conditionBlock);
        loopExit.addSuccessor(joinBlock);

        joinBlock;
      }
    }
  }


  /**
   * Build BasicBlocks for an Expression.
   *
   * @param expression the expression
   * @param currentBlock the basic block to add nodes to
   * @return (the next basic block to add nodes to, the temporary holding the result value)
   */
  def parseExpression(expression: Expression, currentBlock: NormalBlock): (NormalBlock, Int) = {
    expression match {
      case BinaryExpression(id, operator, left, right, location) => {
        val (leftBlock, leftResult) = parseExpression(left, currentBlock);
        val (rightBlock, rightResult) = parseExpression(right, leftBlock);

        val resultVar = tempVar.next();
        val node = BinaryOperatorNode(nodeCounter.next(), operator, leftResult, rightResult, resultVar, location);
        rightBlock.addNode(node);

        (rightBlock, resultVar);
      }
      case IntegerLiteralExpression(id, value, high, location) => {
        val resultVar = tempVar.next();
        val node = IntegerLiteralNode(nodeCounter.next(), value, resultVar, high, location);
        currentBlock.addNode(node);

        (currentBlock, resultVar);
      }
      case ObjectLiteralExpression(id, location) => {
        val resultVar = tempVar.next();
        val node = ObjectLiteralNode(nodeCounter.next(), resultVar, location);
        currentBlock.addNode(node);

        (currentBlock, resultVar);
      }
      case FunctionLiteralExpression(id, body, location) => {
        val function = parseFunction(body);
        val resultVar = tempVar.next();
        val node = FunctionLiteralNode(nodeCounter.next(), function, resultVar, location);
        currentBlock.addNode(node);

        (currentBlock, resultVar);
      }
      case simpleAccess@SimplePropertyAccessExpression(_, _, _) =>
        (currentBlock, parseReadPropertyExpression(simpleAccess, currentBlock));
      case qualifiedAccess@QualifiedPropertyAccessExpression(_, _, _, _) =>
        (currentBlock, parseReadPropertyExpression(qualifiedAccess, currentBlock));
    }
  }

  /**
   * Read a PropertyAccessExpression.
   *
   * @param propertyAccess the property access expression
   * @param currentBlock the basic block to add nodes to
   * @return the temporary holding the result value
   */
  def parseReadPropertyExpression(propertyAccess: PropertyAccessExpression, currentBlock: NormalBlock): Int = {
    propertyAccess match {
      case SimplePropertyAccessExpression(id, name, location) => {
        val resultVar = tempVar.next();
        val node = ReadVariableNode(nodeCounter.next(), name, resultVar, location);
        currentBlock.addNode(node);
        resultVar
      }
      case QualifiedPropertyAccessExpression(id, nestedPropertyAccess, name, location) => {
        val baseVar = parseReadPropertyExpression(nestedPropertyAccess, currentBlock);
        val resultVar = tempVar.next();
        val node = ReadPropertyNode(nodeCounter.next(), name, baseVar, resultVar, location);
        currentBlock.addNode(node);
        resultVar
      }
    }
  }

}

class FlowGraph(val main: Function,
                val functions: Set[Function],
                val functionMap: Map[BasicBlock, Function],
                val predecessorsMap: MultiMap[BasicBlock, BasicBlock],
                val successorsMap: MultiMap[BasicBlock, BasicBlock],
                val ancestorsMap: MultiMap[BasicBlock, BasicBlock],
                val descendantsMap: MultiMap[BasicBlock, BasicBlock]) {

  def functionOf(block: BasicBlock): Function = {
    functionMap.get(block) match {
      case Some(f) => f;
      case None => throw new RuntimeException("Function not found");
    }
  }

  def predecessorsOf(block: BasicBlock): Set[BasicBlock] = predecessorsMap(block);

  def successorsOf(block: BasicBlock): Set[BasicBlock] = successorsMap(block);

  def ancestorsOf(block: BasicBlock): Set[BasicBlock] = ancestorsMap(block);

  def descendantsOf(block: BasicBlock): Set[BasicBlock] = descendantsMap(block);

  def toGraphviz: Set[GraphvizCluster] = {
    functions.map(f => f.toGraphviz);
  }

}